package com.hzsparrow.ai.config;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.Base64;
import java.nio.charset.StandardCharsets;

import org.elasticsearch.client.Request;
import org.elasticsearch.client.Response;
import org.elasticsearch.client.RestClient;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.elasticsearch.ElasticsearchVectorStore;
import org.springframework.ai.vectorstore.elasticsearch.ElasticsearchVectorStoreOptions;
import org.springframework.ai.vectorstore.elasticsearch.SimilarityFunction;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;
import org.apache.http.HttpHost;
import org.apache.http.Header;
import org.apache.http.message.BasicHeader;

import lombok.extern.slf4j.Slf4j;

/**
 * 向量存储管理器
 * 用于动态创建和管理Elasticsearch向量存储
 */
@Slf4j
@Component
public class VectorStoreManager {

    @Value("${spring.elasticsearch.uris:http://localhost:9200}")
    private String elasticsearchUris;
    
    @Value("${spring.elasticsearch.username:}")
    private String elasticsearchUsername;
    
    @Value("${spring.elasticsearch.password:}")
    private String elasticsearchPassword;

    @Value("${spring.ai.vectorstore.elasticsearch.index-name:my-index}")
    private String indexName;
    
    private final Map<String, VectorStore> vectorStoreCache = new ConcurrentHashMap<>();
    
    /**
     * 创建手动配置的RestClient
     */
    private RestClient createRestClient() {
        try {
            log.info("创建手动配置的RestClient: URI={}, 用户名={}", elasticsearchUris, elasticsearchUsername);
            
            // 解析URI
            String uri = elasticsearchUris.split(",")[0]; // 取第一个URI
            HttpHost httpHost = HttpHost.create(uri);
            
            // 创建认证头
            String credentials = elasticsearchUsername + ":" + elasticsearchPassword;
            String encodedCredentials = Base64.getEncoder().encodeToString(
                credentials.getBytes(StandardCharsets.UTF_8));
            
            Header authHeader = new BasicHeader("Authorization", "Basic " + encodedCredentials);
            
            RestClient restClient = RestClient.builder(httpHost)
                .setDefaultHeaders(new Header[]{authHeader})
                .build();
            Response response = restClient.performRequest(new Request("GET", "/"));
            int statusCode = response.getStatusLine().getStatusCode();
            if (statusCode != 200) {
                log.error("创建RestClient失败: 状态码=" + statusCode);
                throw new RuntimeException("创建RestClient失败: 状态码=" + statusCode);
            }
            log.info("RestClient创建成功");
            return restClient;
        } catch (Exception e) {
            log.error("创建RestClient失败: {}", e.getMessage(), e);
            throw new RuntimeException("创建RestClient失败", e);
        }
    }
    
    /**
     * 创建Elasticsearch向量存储
     * 
     * @param embeddingModel 嵌入模型
     * @param indexName 索引名称
     * @param dimensions 向量维度
     * @return 向量存储实例
     */
    public VectorStore createElasticsearchVectorStore(EmbeddingModel embeddingModel, 
                                                     String indexName, 
                                                     int dimensions) {
        try {
            log.info("开始创建Elasticsearch向量存储: 索引={}, 维度={}", indexName, dimensions);
            
            // 创建选项
            ElasticsearchVectorStoreOptions options = new ElasticsearchVectorStoreOptions();
            options.setIndexName(indexName);
            options.setDimensions(dimensions);
            options.setSimilarity(SimilarityFunction.cosine);
            
            // 使用手动配置的RestClient
            RestClient restClient = createRestClient();
            VectorStore vectorStore = ElasticsearchVectorStore.builder(restClient, embeddingModel)
                .options(options)
                .initializeSchema(true)
                .build();
            
            log.info("创建Elasticsearch向量存储成功: {} - 维度: {}", indexName, dimensions);
            return vectorStore;
        } catch (Exception e) {
            log.error("创建Elasticsearch向量存储失败: {} - {}", indexName, e.getMessage(), e);
            throw new RuntimeException("创建Elasticsearch向量存储失败", e);
        }
    }
    
    /**
     * 获取或创建向量存储（带缓存）
     * 
     * @param embeddingModel 嵌入模型
     * @param indexName 索引名称
     * @param dimensions 向量维度
     * @return 向量存储实例
     */
    public VectorStore getOrCreateVectorStore(EmbeddingModel embeddingModel, 
                                             String indexName, 
                                             int dimensions) {
        String cacheKey = indexName + "_" + dimensions;
        return vectorStoreCache.computeIfAbsent(cacheKey, 
            k -> createElasticsearchVectorStore(embeddingModel, indexName, dimensions));
    }
    
    /**
     * 获取默认向量存储
     * 
     * @param embeddingModel 嵌入模型
     * @return 向量存储实例
     */
    public VectorStore getDefaultVectorStore(EmbeddingModel embeddingModel) {
        return getOrCreateVectorStore(embeddingModel, indexName, 1024);
    }
    
    /**
     * 清除指定索引的缓存
     * 
     * @param indexName 索引名称
     * @param dimensions 向量维度
     */
    public void clearCache(String indexName, int dimensions) {
        String cacheKey = indexName + "_" + dimensions;
        vectorStoreCache.remove(cacheKey);
        log.info("清除向量存储缓存: {}", cacheKey);
    }
    
    /**
     * 清除所有缓存
     */
    public void clearAllCache() {
        vectorStoreCache.clear();
        log.info("清除所有向量存储缓存");
    }
}